import cv2
import mmcv
import numpy as np
from mmcv.transforms import Compose
from mmengine.utils import track_iter_progress
from mmdet.registry import VISUALIZERS
from mmdet.apis import init_detector as det_init_detector
from mmdet.apis import inference_detector as det_inference_detector
from glob import glob
from tqdm import tqdm
import sys
import os.path as osp
base_dir = osp.dirname(osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__)))))
sys.path.append("./laketicv")
from ltcv.models.backbones import *
from ltcv.models.necks import *
from ltcv.models.heads import *
from mmdet.apis import DetInferencer
# from ltcv.models.backbones.csp_darknet_deploy import *
def det_init(config_file,checkpoint_file):
    model = det_init_detector(config_file, checkpoint_file, device='cuda:0')
    
    # model = DetInferencer(config_file, checkpoint_file, device='cuda:0')
    return model


# def det_infer(model,img_dirs, pred_score_thr=0.3):
#     # Test a single image and show the results
#     results = []
#     tmp = []
#     predictions = model(img_dirs, batch_size=16, return_vis=True)
#     outputs = predictions['predictions']
#     for i, result in enumerate(outputs):
#         labels = result["labels"]
#         scores = result["scores"]
#         bboxes = result["bboxes"]
#         for label,score,bbox in zip(labels,scores,bboxes):
#             tmp.append((label,score,bbox))
#         results.append((img_dirs[i],tmp))
#     return results

def det_infer(model,img_dirs, pred_score_thr=0.3):
    # Test a single image and show the results
    results = []
    for img_dir in tqdm(img_dirs):
        tmp=[]
        result = det_inference_detector(model, img_dir)
        # print(result)
        # print(type(result))
        scores = result.pred_instances.scores.cpu().numpy()
        index_match = np.where(scores < pred_score_thr)[0][0] if len(np.where(scores < pred_score_thr)[0])>0 else len(scores)
        scores = scores[:index_match]
        labels = result.pred_instances.labels.cpu().numpy()[:index_match]
        bboxes = result.pred_instances.bboxes.cpu().numpy()[:index_match]
        for label,score,bbox in zip(labels,scores,bboxes):
            tmp.append((label,score,list(bbox)))
        results.append((img_dir,tmp))
    return results


# if __name__ == "__main__":
#     det_infer("/workspace/env/config_and_weight/det/yolox_l_1024_4xb12-100e_flip_aoi_weice_copy.py",
#               "/workspace/env/config_and_weight/det/best_coco_bbox_mAP_epoch_99.pth",
#               "/data/dataset/data/weice_data/det/for_test")